# pose_worker.py
from multiprocessing import Process, Event, Queue
import numpy as np
import cv2
import os
import time
import traceback
import queue  # 添加queue模块导入
# from stereo.ppnet.ppnet_rknn_infer_32 import PPnet_Pose_Track_ins, BODY_PARTS
from stereo.ppnet.rtmpose_rknn_infer import PPnet_Pose_Track_ins, BODY_PARTS
class MovingAverageFilter:
    def __init__(self, window_size=5):
        self.window_size = window_size
        self.buffer = []

    def apply(self, keypoints):
        """
        keypoints: numpy array of shape (26, 2)
        returns: smoothed keypoints, shape (26, 2)
        """
        self.buffer.append(keypoints)
        if len(self.buffer) > self.window_size:
            self.buffer.pop(0)
        return np.mean(self.buffer, axis=0)
    
class RealtimeKeypointFilter:
    def __init__(self, alpha=0.6):
        """
        alpha: 平滑系数（越小越平滑，越大越灵敏）
        alpha	描述	推荐场景
        0.2	非常平滑，适合静止	摄像头固定、轻微抖动
        0.3~0.5	通用型，适度响应运动	普通运动姿态估计
        0.6~0.8	快速响应（较少平滑）	快速运动或对延迟敏感
        """
        self.alpha = alpha
        self.prev_kpts = None

    def apply(self, current_kpts):
        """
        current_kpts: 当前帧的关键点，形状为 (26, 2)，float32/float64
        返回: 平滑后的关键点，(26, 2)
        """
        if self.prev_kpts is None:
            self.prev_kpts = current_kpts.copy()
            return current_kpts

        smoothed = self.alpha * current_kpts + (1 - self.alpha) * self.prev_kpts
        self.prev_kpts = smoothed
        return smoothed
RealtimeKeypointFilter_ins = RealtimeKeypointFilter(alpha=0.3)
class PoseEstimationWorker(Process):
    def __init__(self, task_queue=None, result_queue=None, pose_ready_event=None):
        super().__init__()
        # 如果没有传入队列，则创建队列
        self.task_queue = task_queue if task_queue is not None else Queue(maxsize=1)
        self.result_queue = result_queue if result_queue is not None else Queue(maxsize=1)
        self.pose_ready_event = pose_ready_event  # 用于通知主进程姿态估计完成 (multiprocessing.Event)
        
        # 导入PPnet模型
        try:
            self.ppnet = PPnet_Pose_Track_ins
            self.BODY_PARTS = BODY_PARTS
            print("成功加载PPnet姿态估计模型")
            
            # 导入YOLOv6用于检测
            from stereo.yolov6.yolov6 import model_inference
            self.model_inference = model_inference
            print("成功加载YOLOv6检测模型")
            
            self.frame_count = 0  # 帧计数器，用于控制何时执行检测
            
        except ImportError as e:
            print(f"导入姿态估计模型失败: {e}")
            traceback.print_exc()
            self.ppnet = None
            self.model_inference = None

    def run(self):
        """处理队列中的图像帧"""
        print(f"姿态估计进程启动,PID: {os.getpid()}")
        while True:
            try:
                left_image, job_id = self.task_queue.get(timeout=0.5)  # 添加超时避免阻塞
                if left_image is None:
                    time.sleep(0.01)
                    continue
                
                # 姿态估计
                if self.frame_count == 0:
                    # left_image = cv2.resize(left_image, (640, 480))
                    boxes = self.model_inference(left_image)
                    if boxes is None:
                        self.frame_count += 1
                    else:
                        # boxes = None
                        self.frame_count = 0
                    kps_dict_2d, kps_2d_list = self.ppnet(left_image, boxes)
                    self.frame_count += 1
                else:
                    kps_dict_2d, kps_2d_list = self.ppnet(left_image)
                if len(self.ppnet.bboxes_last_frame) > 0:
                    box = self.ppnet.bboxes_last_frame[0]
                    #限制box的x1,y1,x2,y2在图像范围内
                    box = np.clip(box, 0, [640-1, 480-1, 640-1, 480-1])
                # kps_2d_list = RealtimeKeypointFilter_ins.apply(kps_2d_list)
                # kps_2d_list = MovingAverageFilter_ins.apply(kps_2d_list)
                kps_2d_list = kps_2d_list.astype(np.int32)
                # 将结果放入结果队列
                self.result_queue.put((kps_dict_2d, kps_2d_list, box, job_id))
                
                # 设置姿态估计完成事件
                if self.pose_ready_event:
                    self.pose_ready_event.set()
                    
            except queue.Empty:  # 修复引用错误
                # 队列为空，继续等待
                continue
            except Exception as e:
                print(f"姿态估计出错: {e}")
                traceback.print_exc()
                # 即使出错也要设置事件，避免主进程一直等待
                if self.pose_ready_event:
                    self.pose_ready_event.set()
                    

# 使用示例
def create_pose_processor():
    """创建并返回姿态估计处理器及其通信组件"""
    task_queue = Queue(maxsize=1)  # 任务队列
    result_queue = Queue(maxsize=1)  # 结果队列
    pose_ready_event = Event()  # 使用多进程Event
    
    pose_processor = PoseEstimationWorker(
        task_queue=task_queue,
        result_queue=result_queue,
        pose_ready_event=pose_ready_event
    )
    
    # 启动进程
    pose_processor.start()
    
    return pose_processor, task_queue, result_queue, pose_ready_event
